
import torch
from .MCNN import MCNN
from torch.nn.modules.module import Module
from torch.nn.parameter import Parameter
'''
  Res_MCNN网络
  输入 x1（1,n,n）  
  输出（1,n，n）
  示例：
    x1 = torch.rand([100, 1, 69, 69])
    model = Res_MCNN(num_convolutions=3, kernel_size=3, scale=2)
    y = model(x1)
    print(y.shape)
  注意：
    kernel_size必须为单数，否则无法保证每次卷积后，size不变
    O = （I - k + 2p）/S + 1
'''

class Res_MCNN(Module):
    def __init__(self,num_convolutions, kernel_size, scale = 2):
        super(Res_MCNN, self).__init__()
        self.preweight = Parameter(torch.empty(1).fill_(0.5)).type(torch.float32)
        self.nowweight = Parameter(torch.empty(1).fill_(0.5)).type(torch.float32)
        self.kernel_size = kernel_size
        self.num_convolutions = num_convolutions
        self.scale = scale
        self.MCNN = MCNN(num_convolutions=num_convolutions, kernel_size=kernel_size, scale=scale) #MCNN

    def forward(self, x):
        out = self.MCNN(x)
        return self.preweight * x + self.nowweight * out